import os
import numpy as np
import config


def run_with_params(param_dict: dict):
    line = f'python ./src/train.py'
    for name, value in param_dict.items():
        line += f' -{name}={value}'
    print('running', line)
    os.system(line)


def run_with_param_string(pstr: str):
    line = f'python ./src/train.py'
    for param in pstr.split():
        line += f' {param}'
    print(line)
    os.system(line)


def one_param_run(arg_name, arg_list, extra=""):
    for v in arg_list:
        line = f"python ./main.py --{arg_name.replace('_', '_')}={v} {extra}"
        print("running: " + line)
        os.system(line)


def two_param_run(v1_name, v1_list, v2_name, v2_list, extra=""):
    for v1 in v1_list:
        for v2 in v2_list:
            line = f"python ./main.py --{v1_name.replace('_', '_')}={v1} --{v2_name.replace('_', '_')}={v2} {extra}"
            os.system(line)


def power_sir():
    run_with_param_string(
        '''
        -dataset=power
        -infect_rate=0.15
        -recover_rate=0.04
        -batch_size=5
        -seq_len=30
        -data_group_size=30
        -cluster_number=1
        -generate_model=SIR
        '''
    )

def power_si():
    run_with_param_string(
        '''
        -dataset=power
        -infect_rate=0.1
        -recover_rate=0.00
        -batch_size=5
        -seq_len=30
        -num_src=700
        -data_group_size=30
        -cluster_number=1
        -generate_model=SI
        '''
    )





if __name__ == '__main__':
    power_si()
